import numpy as np
from sfepy.discrete.fem import Mesh, FEDomain, Field
from sfepy.discrete import Problem, Functions, Equations, Region
from sfepy.base.base import assert_, output, Struct
from sfepy.solvers import Solver
from sfepy.mechanics.matcoefs import stiffness_from_lame
from sfepy.discrete import (FieldVariable, Material, Integral, Integrals,
                            Equation, Equations, Problem)
from sfepy.mechanics.matcoefs import stiffness_from_youngpoisson
from sfepy.terms import Term
from sfepy.discrete.conditions import Conditions, EssentialBC
import scipy.sparse.linalg as sla
from structAIDetect.simdatagen.getfiles import get_files_by_extension, shape_plot


N_EIGS = 5  #总共分析前5阶模态振型




def modal_analysis(eq1: Equation, eq2: Equation, pb: Problem, cavity, n_eigs, eig_solver):
    # 7. 定义方程
    mtx_k = eq1.evaluate_withcavity(mode='weak', dw_mode='matrix', asm_obj=pb.mtx_a, cavity=cavity)
    mtx_m = mtx_k.copy()
    mtx_m.data[:] = 0.0
    mtx_m = eq2.evaluate_withcavity(mode='weak', dw_mode='matrix', asm_obj=mtx_m, cavity=cavity)
    n_rbm = 0
    try:
        eigs, svecs = eig_solver(mtx_k, mtx_m, n_eigs + n_rbm,
                                    eigenvectors=True)
    except sla.ArpackNoConvergence as ee:
        eigs = ee.eigenvalues
        svecs = ee.eigenvectors
        output('only %d eigenvalues converged!' % len(eigs))

    eigs = eigs[n_rbm:]
    svecs = svecs[:, n_rbm:]
    omegas = np.sqrt(eigs)
    freqs = omegas / (2 * np.pi)
    output('number |         eigenvalue |  angular frequency '
            '|          frequency')
    for ii, eig in enumerate(eigs):
        output('%6d | %17.12e | %17.12e | %17.12e'
                % (ii + 1, eig, omegas[ii], freqs[ii]))

    # 计算振型向量
    variables = pb.set_default_state()

    vecs = np.empty((variables.di.n_dof_total, svecs.shape[1]),
                    dtype=np.float64)
    for ii in range(svecs.shape[1]):
        vecs[:, ii] = variables.make_full_vec(svecs[:, ii])

    return freqs, vecs

def cellinfo_generate(domain: FEDomain, omega: Region, 
                      vecs: np.ndarray, cavity: np.ndarray):
    """
    将结构模态振动分析的结果保存为ndarray的形式。
    
    参数
    ---------
    domain : FEDomain
        分析的几何结构的FEDomain对象。
    vecs : ndarray
        模态分析得到的振型向量。
    omega : Region
        表征domain全局区域的Region，通常通过domain.create_region('Omega', 'all')获取。
    cavity : 
        表示模态分析中刚度和质量被置0的单元索引。
    
    返回值
    ---------
    cellinfo : ndarray
        cellinfo.shape = (cellnum, 8, 19)
        cellinfo的第0维对应domain中的不同单元。
        cellinfo的第1为对应六面体单元中的8个节点。
        cellinfo第2维的0位对应每个单元的节点id。
        cellinfo第2维的1~3位对应每个单元的节点(x,y,z)坐标。
        cellinfo第2维的4~6位对应每个单元的节点在第1阶模态下对应的(x,y,z)位移。
                    7~9位对应每个单元的节点在第2阶模态下对应的(x,y,z)位移。
                    ...
                    16~18位对应每个单元的节点在第5阶模态下对应的(x,y,z)位移。
    
    cellstate : ndarray
        cellstate.shape = (cellnum, )
        每一位表示domain中的每个单元的质量和刚度是否被置0。
    """
    cells_nodes = domain.get_conn()
    node_corr = omega.cmesh.coors
    cells_nodes_corr = node_corr[cells_nodes]
    node_vecs = vecs.reshape(-1, 3 * vecs.shape[1])
    cells_nodes_vec = node_vecs[cells_nodes, :]
    cellinfo = np.zeros([cells_nodes.shape[0], 8, 19], dtype=np.float32)
    cellstate = np.ones(cells_nodes.shape[0], dtype=bool)
    cellinfo[:, :, 0] = cells_nodes
    cellinfo[:, :, 1:4] = cells_nodes_corr
    cellinfo[:, :, 4:19] = cells_nodes_vec
    cellstate[cavity] = False

    return cellinfo, cellstate










